import torch
import torch.nn as nn
import torch.nn.functional as F


class OhemCrossEntyopy(nn.Module):
    def __init__(self, ignore_label=-1, weight=None, thres=0.7, min_kept=100000):
        super(OhemCrossEntyopy, self).__init__()
        self.thres = thres
        self.min_kept = max(1, min_kept)
        self.ignore_label = ignore_label
        self.criterion = nn.CrossEntropyLoss(
            weight=weight,
            ignore_index=ignore_label,
            reduction='none'
        )

    def _ce_forward(self, pred, target):
        ph, pw = pred.size(2), pred.size(3)
        h, w = target.size(1), target.size(2)
        if ph != h or pw != w:
            pred = F.interpolate(input=pred, size=(h, w), mode='bilinear', align_corners=True)

        loss = self.criterion(pred, target)

        return loss

    def _ohem_forward(self, score, target):
        ph, pw = score.size(2), score.size(3)
        h, w = target.size(1), target.size(2)
        if ph != h or pw != w:
            score = F.interpolate(input=score, size=(h, w), mode='bilinear', align_corners=True)
        pred = F.softmax(score, dim=1)
        pixel_losses = self.criterion(score, target).contiguous().view(-1)
        mask = target.contiguous().view(-1) != self.ignore_label

        tmp_target = target.clone()
        tmp_target[tmp_target == self.ignore_label] = 0
        pred = pred.gather(1, tmp_target.unsqueeze(1))
        pred, ind = pred.contiguous().view(-1, )[mask].contiguous().sort()
        min_value = pred[min(self.min_kept, pred.numel() - 1)]
        threshold = max(min_value, self.thresh)

        pixel_losses = pixel_losses[mask][ind]
        pixel_losses = pixel_losses[pred < threshold]
        return pixel_losses.mean()

    def forward(self, pred, target):
        return self._ohem_forward(pred, target)
